# -*- coding: utf-8 -*-
"""
Created on Thu Sep 12 13:55:19 2013

@author: jkwong
"""
# PBAR_BasicScanZspecViewer

import os
import sys
import csv
import time
import matplotlib.cm as cm
import numpy as np
import matplotlib.pyplot as plt

def ReadZspecBasicScan(filePath, prefixList, sliceList):
    """Reads Basic scan zspec files"""
    #2888-FDFC
    dat = []
    energy = []
    NumberDetectors = 136
    NumberEnergyBins = 32
    for index, prefix in enumerate(prefixList): # cycle through datasets
        dat2 = np.zeros((len(sliceList[index]), NumberDetectors, NumberEnergyBins))
        energy2 = np.zeros(NumberEnergyBins)
        count = 0        
        for sliceIndex, slice in enumerate(sliceList[index]):
            if (sliceIndex % 100) == 0:
                print(sliceIndex)
            fullFilename = os.path.join(filePath, prefix + '-slice{}.csv'.format(slice))
            try:
                with open(fullFilename, 'rb') as fid:
                    csvReaderObj = csv.reader(fid)
                    # skip first three lines
                    lineIn = csvReaderObj.next()
                    lineIn = csvReaderObj.next()
                    lineIn = csvReaderObj.next()
                    for i in xrange(NumberEnergyBins):  # cycle over energy bins
                        lineIn = csvReaderObj.next()
                        energy2[i] = float(lineIn[0])
                        dat2[sliceIndex,:,i] = np.array(lineIn[1:-1]).astype(float)
                    count += 1
            except:
                print("No more files")
                break
        # truncate it
        dat2 = dat2[0:count,:,:]
        dat.append(dat2)
        energy.append(energy2)
    return(energy, dat)
    
def ReadZspecBasicScanNumpy(fullFilename):
    """Reads Basic scan zspec files"""
    energy = np.array([7.50E+00,1.05E+01,1.35E+01,1.65E+01,1.95E+01,2.25E+01,2.55E+01, \
        2.85E+01,3.15E+01,3.45E+01,3.75E+01,4.05E+01,4.35E+01,4.65E+01,4.95E+01,\
        5.25E+01,5.55E+01,5.85E+01,6.15E+01,6.45E+01,6.75E+01,7.05E+01,7.35E+01, \
        7.65E+01,7.95E+01,8.25E+01,8.55E+01,8.85E+01,9.15E+01,9.45E+01,9.75E+01,1.01E+02])
    dat = np.load(fullFilename)
    if dat.shape[2] > dat.shape[1]:
        dat = dat.transpose([0,2,1])
    return(energy, dat)

class DisplayZspecBasic:
    def __init__(self, dat, energy):
        self.dat = dat
        self.energy = energy;
        self.fig1 = plt.figure(99)
        self.ax1 = self.fig1.add_subplot(111)
        
        self.ax1.imshow(self.dat.sum(2).T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
        
        self.fig1.canvas.show()
#        self.fig1.show()
        
        (self.x1, self.x2, self.y1, self.y2) = (0,0,0,0)
                
        self.cid1 = self.fig1.canvas.mpl_connect('button_press_event', self.on_press)
        self.cid2 = self.fig1.canvas.mpl_connect('button_release_event', self.on_release)
        self.cid3 = self.fig1.canvas.mpl_connect('motion_notify_event', self.on_motion)
        self.cid4 = self.fig1.canvas.mpl_connect('key_press_event', self.on_key)
        
        self.height = self.dat.shape[1]
        self.width = self.dat.shape[0]
        
        self.fig2 = plt.figure(100)
        self.ax2 = self.fig2.add_subplot(1,1,1)

        self.rect = plt.Rectangle((0,0), 0,0, alpha = 0.5, color = 'g')
        self.fig2Log = 1
        self.fig1Log = 0

    # pressing a key will change the view
    def on_key(self, event):
        print('you pressed', event.key, event.xdata, event.ydata)
        plt.figure(self.fig1.number)
        if (event.key == u'c'):
            temp = self.dat.astype(float)
            intensity = temp.sum(2).T
            plt.imshow(intensity, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
            print("Total Counts")
        elif (event.key == u'm'):  # press m to see mean bin
            energyMatrix = np.tile(self.energy, (self.dat.shape[0], self.dat.shape[1], 1))
            print energyMatrix.shape
            intensity = (self.dat * energyMatrix).sum(2) / self.dat.sum(2)
            plt.imshow(intensity.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
            print("Mean Bin")
        elif (event.key == u'v'):  # v to see standard deviation
            energyMatrix = np.tile(self.energy, (self.dat.shape[0], self.dat.shape[1], 1))
            binMean = (self.dat * energyMatrix).sum(2) / energyMatrix.sum(2)
            binMeanMatrix = np.tile(binMean, (self.dat.shape[0], 1, 1))
            
            intensity = np.sqrt( (( (energyMatrix - binMeanMatrix)**2) * self.dat ).astype(float) / self.dat.sum(2).astype(float)  )
                        
            plt.imshow(intensity.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
            print("Std Bin")
        elif (event.key == u'r'):  # v to see std/mean
            try:
                energyMatrix = np.tile(self.energy, (self.dat.shape[0], self.dat.shape[1], 1))
                binMean = (self.dat * energyMatrix).sum(2) / energyMatrix.sum(2)
                binMeanMatrix = np.tile(binMean, (self.dat.shape[0], 1, 1))
                print binMeanMatrix.shape
                print energyMatrix.shape
#            try:
                binStdMatrix = np.sqrt( (( (energyMatrix - binMeanMatrix)**2) * self.dat ).astype(float) / self.dat.sum(2).astype(float)  )
                print binStdMatrix.shape
            except:
                print('problem')
            intensity = binStdMatrix/binMeanMatrix
            plt.imshow(intensity.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
            print("Std/Mean Bin")
             
#        stats['binMean'][ii,:] = (binArray * datTemp).sum(axis = 0).astype(float) / datTemp.sum(axis = 0).astype(float)
#        temp = np.matlib.repmat(stats['binMean'][ii,:], datTemp.shape[0],1)  # array of means, 256 x number detectors
#    
#        stats['binSTD'][ii,:] = np.sqrt(  ( ((binArray - temp)**2) * datTemp).sum(axis = 0).astype(float) / datTemp.sum(axis = 0).astype(float)  )
#        stats['binSTD_binMean'][ii,:] = stats['binSTD'][ii,:] / stats['binMean'][ii,:]
            
    def on_press(self, event):
        self.x1, self.y1 = event.xdata, event.ydata
        print("Press: ", self.x1, self.y1)
        
    def on_release(self, event):
        self.x2, self.y2 = event.xdata, event.ydata
        print("Release: ", self.x2, self.y2)
        
        x1 = round(self.x1)
        y1 = round(self.y1)          
        x2 = round(self.x2)
        y2 = round(self.y2)
        
        xrange = [min(x1, x2), max(x1, x2)+1]
        yrange = [min(y1, y2), max(y1, y2)+1]

        # highlight the patches
        # first remove the existing patch if it exists
        try:
            if self.ax1.patches !=[]:
                self.ax1.patches.pop()
            self.rect = plt.Rectangle((xrange[0]-0.5,yrange[0]-0.5), xrange[1] - xrange[0], yrange[1] - yrange[0], alpha = 0.5, color = 'g') 
            self.ax1.add_patch(self.rect)
            self.fig1.canvas.draw()
        except:
            print "asdf"

        if (self.x1 == self.x2) and (self.y1 == self.y2):
            if (x1 <= self.width) and (x1 >= 0) and (y1 <= self.height) and (y1 >= 0):
                plt.figure(self.fig2.number)
                plt.clf()

#                self.ax2.plot(self.dat[x1, y1,:])
#                self.ax2.set_title("{}, {}, {}, {}".format(self.x1, self.x2, self.y1, self.y2))
                plt.plot(self.dat[x1, y1,:])
                plt.title("{}, {}, {}, {}".format(self.x1, self.x2, self.y1, self.y2))
                if self.fig2Log:
                    plt.yscale('log')
                else:
                    plt.yscale('linear')
                self.fig2.canvas.show()
        else:
            plt.figure(self.fig2.number)
            plt.clf()
            if (xrange[0] == xrange[1]) and (yrange[0] != yrange[1]):
#                self.ax2.plot(self.dat[xrange[0], yrange[0]:yrange[1],:].mean(0))
                plt.plot(self.dat[xrange[0], yrange[0]:yrange[1],:].mean(0))
                
            elif (xrange[0] != xrange[1]) and (yrange[0] == yrange[1]):
#                self.ax2.plot(self.dat[xrange[0]:xrange[1], yrange[0],:].mean(0))
                plt.plot(self.dat[xrange[0]:xrange[1], yrange[0],:].mean(0))
                
            else:
#                self.ax2.plot(self.dat[xrange[0]:xrange[1], yrange[0]:yrange[1],:].mean(0).mean(0))
                plt.plot(self.dat[xrange[0]:xrange[1], yrange[0]:yrange[1],:].mean(0).mean(0))
            
            plt.title( "x: %3.2f, %3.2f;y: %3.2f, %3.2f\n# Pixels: %d" %(self.x1, self.x2, self.y1, self.y2, (xrange[1] - xrange[0]-1)*(yrange[1] - yrange[0]-1)) )
            if self.fig2Log:
                plt.yscale('log')
            else:
                plt.yscale('linear')
#            self.ax2.set_title("x: %3.2f, %3.2f;y: %3.2f, %3.2f\n# Pixels: %d" \
#                %(self.x1, self.x2, self.y1, self.y2, (len(xrange)-1)*(len(yrange)-1)))
            self.fig2.canvas.show()
             
    def on_motion(self,event):
        1+1
if __name__ == '__main__':
    
    fullfilename = sys.argv[1]
#    filePath, filename = os.path.split(fullfilename)
#    if filePath == '':
#        filePath = '.'
#    prefixList = [filename]
#    sliceList = [np.arange(0,2500)]
#    (energy, dat) = ReadZspecBasicScan(filePath, prefixList, sliceList)
#    c = DisplayZspecBasic(dat[0], energy[0])

    (energy, dat) = ReadZspecBasicScanNumpy(fullfilename)
    
#    while(True):
#        time.sleep(0.1)
